{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import zipfile\n",
    "import subprocess\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET = 'ml-100k'  # only support \"ml-100k\" and \"ml-1m\" now\n",
    "RAW_PATH = os.path.join('./', DATASET)\n",
    "\n",
    "RANDOM_SEED = 0\n",
    "NEG_ITEMS = 99"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Data\n",
    "\n",
    "1. Load interaction data and item metadata\n",
    "2. Filter out items with less than 5 interactions\n",
    "3. Calculate basic statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading data into ./ml-100k\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "100 4808k  100 4808k    0     0  1338k      0  0:00:03  0:00:03 --:--:-- 1337k\n"
     ]
    }
   ],
   "source": [
    "# download data if not exists\n",
    "\n",
    "if not os.path.exists(RAW_PATH):\n",
    "    subprocess.call('mkdir ' + RAW_PATH, shell=True)\n",
    "if not os.path.exists(os.path.join(RAW_PATH, DATASET + '.zip')):\n",
    "    print('Downloading data into ' + RAW_PATH)\n",
    "    subprocess.call(\n",
    "        'cd {} && curl -O http://files.grouplens.org/datasets/movielens/{}.zip'\n",
    "        .format(RAW_PATH, DATASET), shell=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "with zipfile.ZipFile(os.path.join(RAW_PATH, DATASET + '.zip')) as z:\n",
    "    if DATASET == 'ml-100k':\n",
    "        with z.open(os.path.join(DATASET, 'u.data')) as f:\n",
    "            data_df = pd.read_csv(f, sep=\"\\t\", header=None)\n",
    "        with z.open(os.path.join(DATASET, 'u.item')) as f:\n",
    "            meta_df = pd.read_csv(f, sep='|', header=None, encoding='ISO-8859-1')\n",
    "    elif DATASET == 'ml-1m':\n",
    "        with z.open(os.path.join(DATASET, 'ratings.dat')) as f:\n",
    "            data_df = pd.read_csv(f, sep=b'::', header=None, engine='python')\n",
    "        with z.open(os.path.join(DATASET, 'movies.dat')) as f:\n",
    "            meta_df = pd.read_csv(f, sep=b'::', header=None, engine='python')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>user_id</th>\n",
       "      <th>item_id</th>\n",
       "      <th>label</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>196</td>\n",
       "      <td>242</td>\n",
       "      <td>3</td>\n",
       "      <td>881250949</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>186</td>\n",
       "      <td>302</td>\n",
       "      <td>3</td>\n",
       "      <td>891717742</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>22</td>\n",
       "      <td>377</td>\n",
       "      <td>1</td>\n",
       "      <td>878887116</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>244</td>\n",
       "      <td>51</td>\n",
       "      <td>2</td>\n",
       "      <td>880606923</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>166</td>\n",
       "      <td>346</td>\n",
       "      <td>1</td>\n",
       "      <td>886397596</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   user_id  item_id  label       time\n",
       "0      196      242      3  881250949\n",
       "1      186      302      3  891717742\n",
       "2       22      377      1  878887116\n",
       "3      244       51      2  880606923\n",
       "4      166      346      1  886397596"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_df.columns = ['user_id', 'item_id', 'label', 'time']\n",
    "data_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>item_id</th>\n",
       "      <th>i_year</th>\n",
       "      <th>i_Action</th>\n",
       "      <th>i_Adventure</th>\n",
       "      <th>i_Animation</th>\n",
       "      <th>i_Children's</th>\n",
       "      <th>i_Comedy</th>\n",
       "      <th>i_Crime</th>\n",
       "      <th>i_Documentary</th>\n",
       "      <th>i_Drama</th>\n",
       "      <th>...</th>\n",
       "      <th>i_Film-Noir</th>\n",
       "      <th>i_Horror</th>\n",
       "      <th>i_Musical</th>\n",
       "      <th>i_Mystery</th>\n",
       "      <th>i_Romance</th>\n",
       "      <th>i_Sci-Fi</th>\n",
       "      <th>i_Thriller</th>\n",
       "      <th>i_War</th>\n",
       "      <th>i_Western</th>\n",
       "      <th>i_Other</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>01-Jan-1995</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>01-Jan-1995</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>01-Jan-1995</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>01-Jan-1995</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>01-Jan-1995</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 21 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   item_id       i_year  i_Action  i_Adventure  i_Animation  i_Children's  \\\n",
       "0        1  01-Jan-1995         0            0            0             1   \n",
       "1        2  01-Jan-1995         0            1            1             0   \n",
       "2        3  01-Jan-1995         0            0            0             0   \n",
       "3        4  01-Jan-1995         0            1            0             0   \n",
       "4        5  01-Jan-1995         0            0            0             0   \n",
       "\n",
       "   i_Comedy  i_Crime  i_Documentary  i_Drama  ...  i_Film-Noir  i_Horror  \\\n",
       "0         1        1              0        0  ...            0         0   \n",
       "1         0        0              0        0  ...            0         0   \n",
       "2         0        0              0        0  ...            0         0   \n",
       "3         0        1              0        0  ...            0         0   \n",
       "4         0        0              1        0  ...            0         0   \n",
       "\n",
       "   i_Musical  i_Mystery  i_Romance  i_Sci-Fi  i_Thriller  i_War  i_Western  \\\n",
       "0          0          0          0         0           0      0          0   \n",
       "1          0          0          0         0           0      1          0   \n",
       "2          0          0          0         0           0      1          0   \n",
       "3          0          0          0         0           0      0          0   \n",
       "4          0          0          0         0           0      1          0   \n",
       "\n",
       "   i_Other  \n",
       "0        0  \n",
       "1        0  \n",
       "2        0  \n",
       "3        0  \n",
       "4        0  \n",
       "\n",
       "[5 rows x 21 columns]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "genres = [\n",
    "    'i_Action', 'i_Adventure', 'i_Animation', \"i_Children's\", 'i_Comedy', 'i_Crime', \n",
    "    'i_Documentary', 'i_Drama', 'i_Fantasy', 'i_Film-Noir', 'i_Horror', 'i_Musical', \n",
    "    'i_Mystery', 'i_Romance', 'i_Sci-Fi', 'i_Thriller', 'i_War', 'i_Western', 'i_Other'\n",
    "]\n",
    "if DATASET == 'ml-100k':\n",
    "    item_df = meta_df.drop([1, 3, 4], axis=1)\n",
    "    item_df.columns = ['item_id', 'i_year'] + genres\n",
    "elif DATASET == 'ml-1m':\n",
    "    item_df = meta_df.copy()\n",
    "    item_df.columns = ['item_id', 'title', 'genre']\n",
    "    item_df['title'] = item_df['title'].apply(lambda x: x.decode('ISO-8859-1'))\n",
    "    item_df['genre'] = item_df['genre'].apply(lambda x: x.decode('ISO-8859-1'))\n",
    "    genre_dict = dict()\n",
    "    for g in genres:\n",
    "        genre_dict[g] = np.zeros(len(item_df), dtype=np.int32)\n",
    "    item_genre = item_df['genre'].apply(lambda x: x.split('|')).values\n",
    "    for idx, genre_lst in enumerate(item_genre):\n",
    "        for g in genre_lst:\n",
    "            genre_dict['i_' + g][idx] = 1\n",
    "    for g in genres:\n",
    "        item_df[g] = genre_dict[g]\n",
    "    item_df = item_df.drop(columns=['genre'])\n",
    "item_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Filter before: 100000\n",
      "Filter after: 99287\n"
     ]
    }
   ],
   "source": [
    "# Only retain users and items with at least 5 associated interactions\n",
    "\n",
    "print('Filter before:', len(data_df))\n",
    "filter_before = -1\n",
    "while filter_before != len(data_df):\n",
    "    filter_before = len(data_df)\n",
    "    for stage in ['user_id', 'item_id']:\n",
    "        val_cnt = data_df[stage].value_counts()\n",
    "        cnt_df = pd.DataFrame({stage: val_cnt.index, 'cnt': val_cnt.values})\n",
    "        data_df = pd.merge(data_df, cnt_df, on=stage, how='left')\n",
    "        data_df = data_df[data_df['cnt'] >= 5].drop(columns=['cnt'])\n",
    "print('Filter after:', len(data_df))\n",
    "\n",
    "item_df = item_df[item_df['item_id'].isin(data_df['item_id'])]  # remove unuseful metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_users = data_df['user_id'].value_counts().size\n",
    "n_items = data_df['item_id'].value_counts().size\n",
    "n_clicks = len(data_df)\n",
    "min_time = data_df['time'].min()\n",
    "max_time = data_df['time'].max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Users: 943\n",
      "# Items: 1349\n",
      "# Interactions: 99287\n",
      "Time Span: 1997-09-20/1998-04-22\n"
     ]
    }
   ],
   "source": [
    "time_format = '%Y-%m-%d'\n",
    "\n",
    "print('# Users:', n_users)\n",
    "print('# Items:', n_items)\n",
    "print('# Interactions:', n_clicks)\n",
    "print('Time Span: {}/{}'.format(\n",
    "    datetime.utcfromtimestamp(min_time).strftime(time_format),\n",
    "    datetime.utcfromtimestamp(max_time).strftime(time_format))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build Dataset\n",
    "\n",
    "### Interaction data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(RANDOM_SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>user_id</th>\n",
       "      <th>item_id</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>259</td>\n",
       "      <td>255</td>\n",
       "      <td>874724710</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>259</td>\n",
       "      <td>286</td>\n",
       "      <td>874724727</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>259</td>\n",
       "      <td>298</td>\n",
       "      <td>874724754</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>259</td>\n",
       "      <td>185</td>\n",
       "      <td>874724781</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>259</td>\n",
       "      <td>173</td>\n",
       "      <td>874724843</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   user_id  item_id       time\n",
       "0      259      255  874724710\n",
       "1      259      286  874724727\n",
       "2      259      298  874724754\n",
       "3      259      185  874724781\n",
       "4      259      173  874724843"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out_df = data_df[['user_id', 'item_id', 'time']]\n",
    "out_df = out_df.drop_duplicates(['user_id', 'item_id', 'time'])\n",
    "out_df.sort_values(by=['time', 'user_id'], kind='mergesort', inplace=True)\n",
    "out_df = out_df.reset_index(drop=True)\n",
    "out_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>user_id</th>\n",
       "      <th>item_id</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>259</td>\n",
       "      <td>254</td>\n",
       "      <td>874724710</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>259</td>\n",
       "      <td>285</td>\n",
       "      <td>874724727</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>259</td>\n",
       "      <td>297</td>\n",
       "      <td>874724754</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>259</td>\n",
       "      <td>184</td>\n",
       "      <td>874724781</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>259</td>\n",
       "      <td>172</td>\n",
       "      <td>874724843</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   user_id  item_id       time\n",
       "0      259      254  874724710\n",
       "1      259      285  874724727\n",
       "2      259      297  874724754\n",
       "3      259      184  874724781\n",
       "4      259      172  874724843"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# reindex (start from 1)\n",
    "\n",
    "uids = sorted(out_df['user_id'].unique())\n",
    "user2id = dict(zip(uids, range(1, len(uids) + 1)))\n",
    "iids = sorted(out_df['item_id'].unique())\n",
    "item2id = dict(zip(iids, range(1, len(iids) + 1)))\n",
    "\n",
    "out_df['user_id'] = out_df['user_id'].apply(lambda x: user2id[x])\n",
    "out_df['item_id'] = out_df['item_id'].apply(lambda x: item2id[x])\n",
    "out_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# leave one out spliting\n",
    "\n",
    "clicked_item_set = dict()\n",
    "for user_id, seq_df in out_df.groupby('user_id'):\n",
    "    clicked_item_set[user_id] = set(seq_df['item_id'].values.tolist())\n",
    "    \n",
    "def generate_dev_test(data_df):\n",
    "    result_dfs = []\n",
    "    n_items = data_df['item_id'].value_counts().size\n",
    "    for idx in range(2):\n",
    "        result_df = data_df.groupby('user_id').tail(1).copy()\n",
    "        data_df = data_df.drop(result_df.index)\n",
    "        neg_items = np.random.randint(1, n_items + 1, (len(result_df), NEG_ITEMS))\n",
    "        for i, uid in enumerate(result_df['user_id'].values):\n",
    "            user_clicked = clicked_item_set[uid]\n",
    "            for j in range(len(neg_items[i])):\n",
    "                while neg_items[i][j] in user_clicked:\n",
    "                    neg_items[i][j] = np.random.randint(1, n_items + 1)\n",
    "        result_df['neg_items'] = neg_items.tolist()\n",
    "        result_dfs.append(result_df)\n",
    "    return result_dfs, data_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(97401, 943, 943)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "leave_df = out_df.groupby('user_id').head(1)\n",
    "data_df = out_df.drop(leave_df.index)\n",
    "\n",
    "[test_df, dev_df], data_df = generate_dev_test(data_df)\n",
    "train_df = pd.concat([leave_df, data_df]).sort_index()\n",
    "\n",
    "len(train_df), len(dev_df), len(test_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>user_id</th>\n",
       "      <th>item_id</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>259</td>\n",
       "      <td>254</td>\n",
       "      <td>874724710</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>259</td>\n",
       "      <td>285</td>\n",
       "      <td>874724727</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>259</td>\n",
       "      <td>297</td>\n",
       "      <td>874724754</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>259</td>\n",
       "      <td>184</td>\n",
       "      <td>874724781</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>259</td>\n",
       "      <td>172</td>\n",
       "      <td>874724843</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   user_id  item_id       time\n",
       "0      259      254  874724710\n",
       "1      259      285  874724727\n",
       "2      259      297  874724754\n",
       "3      259      184  874724781\n",
       "4      259      172  874724843"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>user_id</th>\n",
       "      <th>item_id</th>\n",
       "      <th>time</th>\n",
       "      <th>neg_items</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>867</th>\n",
       "      <td>821</td>\n",
       "      <td>71</td>\n",
       "      <td>874793969</td>\n",
       "      <td>[685, 560, 1217, 836, 764, 1034, 278, 600, 109...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>987</th>\n",
       "      <td>817</td>\n",
       "      <td>814</td>\n",
       "      <td>874816007</td>\n",
       "      <td>[369, 918, 202, 384, 371, 556, 955, 861, 131, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1126</th>\n",
       "      <td>893</td>\n",
       "      <td>761</td>\n",
       "      <td>874830424</td>\n",
       "      <td>[702, 728, 556, 33, 1036, 213, 1117, 624, 771,...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1801</th>\n",
       "      <td>933</td>\n",
       "      <td>1172</td>\n",
       "      <td>874939247</td>\n",
       "      <td>[1112, 161, 148, 585, 456, 1112, 1038, 1339, 1...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2098</th>\n",
       "      <td>21</td>\n",
       "      <td>759</td>\n",
       "      <td>874951916</td>\n",
       "      <td>[127, 578, 890, 1349, 977, 113, 651, 238, 991,...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      user_id  item_id       time  \\\n",
       "867       821       71  874793969   \n",
       "987       817      814  874816007   \n",
       "1126      893      761  874830424   \n",
       "1801      933     1172  874939247   \n",
       "2098       21      759  874951916   \n",
       "\n",
       "                                              neg_items  \n",
       "867   [685, 560, 1217, 836, 764, 1034, 278, 600, 109...  \n",
       "987   [369, 918, 202, 384, 371, 556, 955, 861, 131, ...  \n",
       "1126  [702, 728, 556, 33, 1036, 213, 1117, 624, 771,...  \n",
       "1801  [1112, 161, 148, 585, 456, 1112, 1038, 1339, 1...  \n",
       "2098  [127, 578, 890, 1349, 977, 113, 651, 238, 991,...  "
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save results\n",
    "\n",
    "train_df.to_csv(os.path.join(RAW_PATH, 'train.csv'), sep='\\t', index=False)\n",
    "dev_df.to_csv(os.path.join(RAW_PATH, 'dev.csv'), sep='\\t', index=False)\n",
    "test_df.to_csv(os.path.join(RAW_PATH, 'test.csv'), sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Item Metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>item_id</th>\n",
       "      <th>i_year</th>\n",
       "      <th>i_Action</th>\n",
       "      <th>i_Adventure</th>\n",
       "      <th>i_Animation</th>\n",
       "      <th>i_Children's</th>\n",
       "      <th>i_Comedy</th>\n",
       "      <th>i_Crime</th>\n",
       "      <th>i_Documentary</th>\n",
       "      <th>i_Drama</th>\n",
       "      <th>...</th>\n",
       "      <th>i_Film-Noir</th>\n",
       "      <th>i_Horror</th>\n",
       "      <th>i_Musical</th>\n",
       "      <th>i_Mystery</th>\n",
       "      <th>i_Romance</th>\n",
       "      <th>i_Sci-Fi</th>\n",
       "      <th>i_Thriller</th>\n",
       "      <th>i_War</th>\n",
       "      <th>i_Western</th>\n",
       "      <th>i_Other</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>13</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>13</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>13</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>13</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>13</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 21 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   item_id  i_year  i_Action  i_Adventure  i_Animation  i_Children's  \\\n",
       "0        1      13         0            0            0             1   \n",
       "1        2      13         0            1            1             0   \n",
       "2        3      13         0            0            0             0   \n",
       "3        4      13         0            1            0             0   \n",
       "4        5      13         0            0            0             0   \n",
       "\n",
       "   i_Comedy  i_Crime  i_Documentary  i_Drama  ...  i_Film-Noir  i_Horror  \\\n",
       "0         1        1              0        0  ...            0         0   \n",
       "1         0        0              0        0  ...            0         0   \n",
       "2         0        0              0        0  ...            0         0   \n",
       "3         0        1              0        0  ...            0         0   \n",
       "4         0        0              1        0  ...            0         0   \n",
       "\n",
       "   i_Musical  i_Mystery  i_Romance  i_Sci-Fi  i_Thriller  i_War  i_Western  \\\n",
       "0          0          0          0         0           0      0          0   \n",
       "1          0          0          0         0           0      1          0   \n",
       "2          0          0          0         0           0      1          0   \n",
       "3          0          0          0         0           0      0          0   \n",
       "4          0          0          0         0           0      1          0   \n",
       "\n",
       "   i_Other  \n",
       "0        0  \n",
       "1        0  \n",
       "2        0  \n",
       "3        0  \n",
       "4        0  \n",
       "\n",
       "[5 rows x 21 columns]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "item_df['item_id'] = item_df['item_id'].apply(lambda x: item2id[x])\n",
    "\n",
    "if DATASET == 'ml-1m':\n",
    "    item_df['i_year'] = item_df['title'].apply(lambda x: int(re.match('.+\\((\\d{4})\\)$', x).group(1)))\n",
    "    item_df = item_df.drop(columns=['title'])\n",
    "elif DATASET == 'ml-100k':\n",
    "    item_df['i_year'] = item_df['i_year'].apply(lambda x: int(str(x).split('-')[-1]) if pd.notnull(x) else 0)\n",
    "seps = [1900, 1940, 1950, 1960, 1970, 1980, 1985] + list(range(1990, int(item_df['i_year'].max() + 2)))\n",
    "year_dict = {}\n",
    "for i, sep in enumerate(seps[:-1]):\n",
    "    for j in range(seps[i], seps[i + 1]):\n",
    "        year_dict[j] = i + 1\n",
    "item_df['i_year'] = item_df['i_year'].apply(lambda x: year_dict[x] if x > 0 else 0)\n",
    "    \n",
    "item_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save results\n",
    "\n",
    "item_df.to_csv(os.path.join(RAW_PATH, 'item_meta.csv'), sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>item_id</th>\n",
       "      <th>title</th>\n",
       "      <th>i_year</th>\n",
       "      <th>i_Action</th>\n",
       "      <th>i_Adventure</th>\n",
       "      <th>i_Animation</th>\n",
       "      <th>i_Children's</th>\n",
       "      <th>i_Comedy</th>\n",
       "      <th>i_Crime</th>\n",
       "      <th>i_Documentary</th>\n",
       "      <th>...</th>\n",
       "      <th>i_Film-Noir</th>\n",
       "      <th>i_Horror</th>\n",
       "      <th>i_Musical</th>\n",
       "      <th>i_Mystery</th>\n",
       "      <th>i_Romance</th>\n",
       "      <th>i_Sci-Fi</th>\n",
       "      <th>i_Thriller</th>\n",
       "      <th>i_War</th>\n",
       "      <th>i_Western</th>\n",
       "      <th>i_Other</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>Toy Story (1995)</td>\n",
       "      <td>01-Jan-1995</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>GoldenEye (1995)</td>\n",
       "      <td>01-Jan-1995</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>Four Rooms (1995)</td>\n",
       "      <td>01-Jan-1995</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>Get Shorty (1995)</td>\n",
       "      <td>01-Jan-1995</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>Copycat (1995)</td>\n",
       "      <td>01-Jan-1995</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 22 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   item_id              title       i_year  i_Action  i_Adventure  \\\n",
       "0        1   Toy Story (1995)  01-Jan-1995         0            0   \n",
       "1        2   GoldenEye (1995)  01-Jan-1995         0            1   \n",
       "2        3  Four Rooms (1995)  01-Jan-1995         0            0   \n",
       "3        4  Get Shorty (1995)  01-Jan-1995         0            1   \n",
       "4        5     Copycat (1995)  01-Jan-1995         0            0   \n",
       "\n",
       "   i_Animation  i_Children's  i_Comedy  i_Crime  i_Documentary  ...  \\\n",
       "0            0             1         1        1              0  ...   \n",
       "1            1             0         0        0              0  ...   \n",
       "2            0             0         0        0              0  ...   \n",
       "3            0             0         0        1              0  ...   \n",
       "4            0             0         0        0              1  ...   \n",
       "\n",
       "   i_Film-Noir  i_Horror  i_Musical  i_Mystery  i_Romance  i_Sci-Fi  \\\n",
       "0            0         0          0          0          0         0   \n",
       "1            0         0          0          0          0         0   \n",
       "2            0         0          0          0          0         0   \n",
       "3            0         0          0          0          0         0   \n",
       "4            0         0          0          0          0         0   \n",
       "\n",
       "   i_Thriller  i_War  i_Western  i_Other  \n",
       "0           0      0          0        0  \n",
       "1           0      1          0        0  \n",
       "2           0      1          0        0  \n",
       "3           0      0          0        0  \n",
       "4           0      1          0        0  \n",
       "\n",
       "[5 rows x 22 columns]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get text\n",
    "genres = [\n",
    "    'i_Action', 'i_Adventure', 'i_Animation', \"i_Children's\", 'i_Comedy', 'i_Crime', \n",
    "    'i_Documentary', 'i_Drama', 'i_Fantasy', 'i_Film-Noir', 'i_Horror', 'i_Musical', \n",
    "    'i_Mystery', 'i_Romance', 'i_Sci-Fi', 'i_Thriller', 'i_War', 'i_Western', 'i_Other'\n",
    "]\n",
    "if DATASET == 'ml-100k':\n",
    "    item_df = meta_df.drop([3, 4], axis=1)\n",
    "    item_df.columns = ['item_id', 'title', 'i_year'] + genres\n",
    "elif DATASET == 'ml-1m':\n",
    "    item_df = meta_df.copy()\n",
    "    item_df.columns = ['item_id', 'title', 'genre']\n",
    "    item_df['title'] = item_df['title'].apply(lambda x: x.decode('ISO-8859-1'))\n",
    "    item_df['genre'] = item_df['genre'].apply(lambda x: x.decode('ISO-8859-1'))\n",
    "    genre_dict = dict()\n",
    "    for g in genres:\n",
    "        genre_dict[g] = np.zeros(len(item_df), dtype=np.int32)\n",
    "    item_genre = item_df['genre'].apply(lambda x: x.split('|')).values\n",
    "    for idx, genre_lst in enumerate(item_genre):\n",
    "        for g in genre_lst:\n",
    "            genre_dict['i_' + g][idx] = 1\n",
    "    for g in genres:\n",
    "        item_df[g] = genre_dict[g]\n",
    "    item_df = item_df.drop(columns=['genre'])\n",
    "item_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 44%|████▎     | 735/1682 [00:00<00:00, 7348.48it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1682/1682 [00:00<00:00, 9834.84it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "outf = open(os.path.join(RAW_PATH, 'ml-100k.text'), 'w')\n",
    "for index, row in tqdm(item_df.iterrows(),total=len(item_df)):\n",
    "    title = row['title']\n",
    "    item_id = row['item_id']\n",
    "    if item_id in item2id:\n",
    "        for g in genres:\n",
    "            if row[g] == 1:\n",
    "                title += ' ' + g[2:]\n",
    "        outf.write(str(item_id) + '\\t' + title + '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_unit2index(file):\n",
    "    unit2index = dict()\n",
    "    with open(file, 'r') as fp:\n",
    "        for line in fp:\n",
    "            unit, index = line.strip().split('\\t')\n",
    "            unit2index[unit] = int(index)\n",
    "    return unit2index\n",
    "\n",
    "\n",
    "def write_remap_index(unit2index, file):\n",
    "    with open(file, 'w') as fp:\n",
    "        for unit in unit2index:\n",
    "            fp.write(str(unit) + '\\t' + str(unit2index[unit]) + '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "write_remap_index(user2id, os.path.join(RAW_PATH, f'ml-100k.user2index'))\n",
    "write_remap_index(item2id, os.path.join(RAW_PATH, f'ml-100k.item2index'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "prompt",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
